﻿Option Explicit On

Imports System.Xml

''' <summary>
''' Contains functions for processing disks from Environment.xml
''' </summary>
''' <remarks></remarks>
Public Class Disks


    ''' <summary>
    ''' High level disk processing, can create, format, 
    ''' </summary>
    ''' <param name="Disk"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Shared Function Process(ByRef Disk As XmlElement) As Boolean

        'checks if the disk exists

        'get drive details
        Dim sLetter As String = Disk.GetAttribute("disk-letter")
        Dim oDrive As IO.DriveInfo = DriveExists(sLetter)
        If oDrive Is Nothing AndAlso CanCreateDrive(Disk) Then


            'try to attach
            If Disk.GetAttribute("ebs-id") <> "" AndAlso Not AttachEBS(Disk) Then Return False

            'try to find something else if couldn;t attach
            Dim blnAttached As Boolean = FindMatchingEBS(Disk)
            If blnAttached AndAlso Not AttachEBS(Disk) Then blnAttached = False

            'try to create if not already attached
            If Not blnAttached AndAlso Disk.GetAttribute("ebs-size") <> "" AndAlso CreateEBS(Disk) Then blnAttached = AttachEBS(Disk)

            'any luck?
            If Not blnAttached Then Return False


            'An EBS was attached - check if there is a new disk
            Threading.Thread.CurrentThread.Sleep(3000)
            Dim iDiskCount As Integer = WmiDisk.WmiLoadDrives

            'bring all disks online
            For i As Integer = 1 To iDiskCount - 1
                DiskPartOnline(i + 1)
            Next
            Threading.Thread.CurrentThread.Sleep(3000)
            iDiskCount = WmiDisk.WmiLoadDrives

            'create partitions
            If Not CreatePartitions() Then Return False
            Threading.Thread.CurrentThread.Sleep(3000)
            iDiskCount = WmiDisk.WmiLoadDrives

            'create volumes
            If Not WmiDisk.WmiFormatVolumes Then Return False
            Threading.Thread.CurrentThread.Sleep(3000)

            Dim iLetterDisk As Integer = WmiDisk.WmiUnnamedVolumeDiskIndex
            If iLetterDisk > 0 AndAlso Not DiskPartLetter(iLetterDisk, sLetter.Chars(0)) Then
                EnviroBuilder.WriteToLog("Cannot set letter " + sLetter + " to disk " + iLetterDisk.ToString, EventLogEntryType.Error)
                Return False
            End If

            'we should have the drive available by letter now
            oDrive = DriveExists(sLetter)

            'must restart if any new disks were added
            EnviroBuilder.MustRestart = True

        End If




        'do we have a drive in the end?
        If oDrive Is Nothing Then
            Return False
        End If

        'is the drive ready to be used?
        If Not oDrive.IsReady Then Return False

        'the drive is fine
        Return True

    End Function


    ''' <summary>
    ''' A safe way to get the drive info
    ''' </summary>
    ''' <param name="Letter"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function DriveExists(ByVal Letter As String) As IO.DriveInfo

        Dim oDrive As IO.DriveInfo = Nothing
        Try
            oDrive = New IO.DriveInfo(Letter)
        Catch ex As Exception
            Return Nothing
        End Try

        If oDrive Is Nothing OrElse Not oDrive.IsReady Then Return Nothing

        Return oDrive

    End Function


    ''' <summary>
    ''' Returns true if the drive can be created from an EBS
    ''' </summary>
    ''' <param name="Disk"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function CanCreateDrive(ByRef Disk As XmlElement) As Boolean

        If Disk.GetAttribute("ebs-id") <> "" Then Return True
        If Disk.GetAttribute("ebs-size") <> "" Then Return True

        Return False

    End Function


#Region " EBS management "

    ''' <summary>
    ''' Creates a new EBS as per the spec and returns true when it was created, which may take a while.
    ''' </summary>
    ''' <param name="xDisk"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function CreateEBS(ByRef xDisk As XmlElement) As Boolean

        'prepare a request
        Dim oClient As New Amazon.EC2.AmazonEC2Client(EnviroBuilder.sAwsKey, EnviroBuilder.sAwsKeySecret)
        Dim oReq As New Amazon.EC2.Model.CreateVolumeRequest
        oReq.AvailabilityZone = EnviroBuilder.oInstance.RunningInstance(0).Placement.AvailabilityZone
        oReq.Size = xDisk.GetAttribute("ebs-size")
        oReq.SnapshotId = xDisk.GetAttribute("ebs-snap")

        EnviroBuilder.WriteToLog("Creating EBS " + oReq.Size + "/" + oReq.SnapshotId, EventLogEntryType.Information)

        Dim sVolID As String = Nothing

        Try
            Dim oResp As Amazon.EC2.Model.CreateVolumeResponse = oClient.CreateVolume(oReq)
            sVolID = oResp.CreateVolumeResult.Volume.VolumeId
        Catch ex As Exception
            EnviroBuilder.WriteToLog("Cannot create EBS" + vbCrLf + ex.Message, EventLogEntryType.Error)
            Return False
        End Try

        If sVolID = "" Then Return False
        xDisk.SetAttribute("ebs-id", sVolID)
        EnviroBuilder.WriteToLog("Created EBS #" + sVolID, EventLogEntryType.Information)

        'wait for it to be created
        Dim oVolReq As New Amazon.EC2.Model.DescribeVolumesRequest
        oVolReq.VolumeId.Add(sVolID)

        For i As Integer = 0 To 60
            Dim oVolResp As Amazon.EC2.Model.DescribeVolumesResponse = Nothing
            Try
                oVolResp = oClient.DescribeVolumes(oVolReq)
                Dim sStatus As String = oVolResp.DescribeVolumesResult.Volume(0).Status
                If sStatus = "available" Then Return True
            Catch ex As Exception
                EnviroBuilder.WriteToLog("EBS status not available" + vbCrLf + ex.Message, EventLogEntryType.Warning)
            End Try

            Threading.Thread.CurrentThread.Sleep(1000)

        Next

        'could not access the new EBS
        EnviroBuilder.WriteToLog("A new EBS is not available: " + sVolID, EventLogEntryType.Error)
        xDisk.RemoveAttribute("ebs-id")

        Return False

    End Function


    ''' <summary>
    ''' Attach an EBS if not attached yet.
    ''' </summary>
    ''' <param name="xDisk"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function AttachEBS(ByRef xDisk As XmlElement) As Boolean

        'volume ID to work with
        Dim sVolID As String = xDisk.GetAttribute("ebs-id")

        'check current status
        Dim iStatus As enAttachmentStatus = VolumeStatus(sVolID)
        EnviroBuilder.WriteToLog("EBS #" + sVolID + " initial status " + iStatus.ToString, EventLogEntryType.Information)
        If iStatus = enAttachmentStatus.Attached Then Return True
        If iStatus = enAttachmentStatus.CannotAttach Then Return False

        'find an available device
        Dim sDeviceID As String = AvailableInstanceDevice()
        If sDeviceID = "" Then Return False

        'should be all good to attach
        Dim oClient As New Amazon.EC2.AmazonEC2Client(EnviroBuilder.sAwsKey, EnviroBuilder.sAwsKeySecret)
        Dim oReq As New Amazon.EC2.Model.AttachVolumeRequest
        oReq.InstanceId = EnviroBuilder.sInstanceID
        oReq.VolumeId = sVolID
        oReq.Device = sDeviceID
        Dim oResp As Amazon.EC2.Model.AttachVolumeResponse = Nothing

        'try to attach
        Try
            oResp = oClient.AttachVolume(oReq)
        Catch ex As Exception
            EnviroBuilder.WriteToLog("Cannot attach volume " + sVolID + " to " + EnviroBuilder.sInstanceID + vbCrLf + ex.Message, EventLogEntryType.Error)
            Return False
        End Try
        'update status
        iStatus = VolumeStatus(sVolID)
        EnviroBuilder.WriteToLog("EBS #" + sVolID + " status changed to " + iStatus.ToString, EventLogEntryType.Information)
        If iStatus = enAttachmentStatus.Attached Then Return True

        'should have been attached by now
        Return False

    End Function


    ''' <summary>
    ''' Gives the fist available device name on the instance from xdva-xdvl
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function AvailableInstanceDevice() As String

        Dim arDeviceName() As String = {"xvdd", "xvde", "xvdf", "xvdg", "xvdh", "xvdi", "xvdj", "xvdk", "xvdl"}

        Dim oClient As New Amazon.EC2.AmazonEC2Client(EnviroBuilder.sAwsKey, EnviroBuilder.sAwsKeySecret)
        Dim oVolReq As New Amazon.EC2.Model.DescribeVolumesRequest

        'get volume description
        Dim oVolResp As Amazon.EC2.Model.DescribeVolumesResponse = Nothing
        For i As Integer = 0 To 5
            Try
                oVolResp = oClient.DescribeVolumes(oVolReq)
                Exit For
            Catch ex As Exception
                EnviroBuilder.WriteToLog("Cannot get list of volumes" + vbCrLf + ex.Message, EventLogEntryType.Error)
            End Try
            Return ""
        Next

        'check status and instance ID
        For Each oVolume As Amazon.EC2.Model.Volume In oVolResp.DescribeVolumesResult.Volume

            For Each oAttachement As Amazon.EC2.Model.Attachment In oVolume.Attachment
                'attached here?
                Dim sInstanceID As String = oAttachement.InstanceId
                If sInstanceID <> EnviroBuilder.sInstanceID Then Continue For
                Dim sDeviceID As String = oAttachement.Device
                'remove a busy device
                If sDeviceID <> "" Then
                    For i As Integer = 0 To arDeviceName.Length - 1
                        If arDeviceName(i).ToLower = sDeviceID Then
                            arDeviceName(i) = ""
                            Exit For
                        End If
                    Next
                End If
            Next
        Next

        'find the first unused device
        For i As Integer = 0 To arDeviceName.Length - 1
            If arDeviceName(i) <> "" Then Return arDeviceName(i)
        Next

        EnviroBuilder.WriteToLog("Instance has no available devices to connect an EBS.", EventLogEntryType.Warning)

        'all are busy
        Return Nothing

    End Function


    ''' <summary>
    ''' Finds a matching EBS by snap-id and then by size. Sets @ebs-id if found
    ''' </summary>
    ''' <param name="xDisk"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function FindMatchingEBS(ByRef xDisk As XmlElement) As Boolean

        Dim sSize As String = xDisk.GetAttribute("ebs-size")
        Dim sSnapID As String = xDisk.GetAttribute("ebs-snap")
        If sSize = "" And sSnapID = "" Then Return False

        Dim sZone As String = EnviroBuilder.oInstance.RunningInstance(0).Placement.AvailabilityZone

        Dim oClient As New Amazon.EC2.AmazonEC2Client(EnviroBuilder.sAwsKey, EnviroBuilder.sAwsKeySecret)
        Dim oVolReq As New Amazon.EC2.Model.DescribeVolumesRequest

        'get list of volumes
        Dim oVolResp As Amazon.EC2.Model.DescribeVolumesResponse = Nothing
        Try
            oVolResp = oClient.DescribeVolumes(oVolReq)
        Catch ex As Exception
            EnviroBuilder.WriteToLog("Cannot get list of EBS." + vbCrLf + ex.Message, EventLogEntryType.Error)
            Return False
        End Try

        Dim sVolID As String = Nothing

        'try to match the volume on snapsho ID
        If sSnapID <> "" Then
            For Each oVolume As Amazon.EC2.Model.Volume In oVolResp.DescribeVolumesResult.Volume
                If oVolume.AvailabilityZone <> sZone Then Continue For
                If oVolume.SnapshotId <> sSnapID Then Continue For
                If oVolume.Status <> "available" Then Continue For

                sVolID = oVolume.VolumeId
                EnviroBuilder.WriteToLog("Found matching EBS on snapshot: " + sVolID, EventLogEntryType.Information)
                Exit For
            Next
        End If

        'do we need to match on size?
        If sVolID = "" AndAlso sSize <> "" Then
            For Each oVolume As Amazon.EC2.Model.Volume In oVolResp.DescribeVolumesResult.Volume
                If oVolume.AvailabilityZone <> sZone Then Continue For
                If oVolume.Status <> "available" Then Continue For
                If oVolume.Size <> sSize Then Continue For

                sVolID = oVolume.VolumeId
                EnviroBuilder.WriteToLog("Found matching EBS on size: " + sVolID, EventLogEntryType.Information)
                Exit For
            Next
        End If

        'volume found?
        If sVolID = "" Then Return False

        xDisk.SetAttribute("ebs-id", sVolID)
        Return True

    End Function

#End Region


#Region " DISKPART "


    ''' <summary>
    ''' Executes diskpart to bring a disk online and returns true if no failures occured
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function DiskPartOnline(ByVal DiskNum As Integer) As Boolean

        'generate file names
        Dim sFileName As String = IO.Path.GetTempFileName
        Dim sLog As String = IO.Path.GetTempFileName

        'save commands file
        Dim oFile As New IO.StreamWriter(sFileName)
        oFile.WriteLine("rescan")
        oFile.WriteLine("select disk " + DiskNum.ToString)
        oFile.WriteLine("online disk")
        oFile.Flush()
        oFile.Close()

        'produce the command line string
        Dim sCmd As String = "diskpart /S """ + sFileName + """" ' > """ + sLog + """"

        'execute
        Dim iRes As Integer = Shell(sCmd, AppWinStyle.Hide, True, 30000)

        'return results
        If iRes <> 0 Then
            EnviroBuilder.WriteToLog("Diskpart ONLINE failed", EventLogEntryType.Warning)
            Return False
        End If
        Return True

    End Function



    ''' <summary>
    ''' Executes diskpart to bring a disk online and returns true if no failures occured
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function DiskPartPartition(ByVal DiskNum As Integer) As Boolean

        'generate file names
        Dim sFileName As String = IO.Path.GetTempFileName
        Dim sLog As String = IO.Path.GetTempFileName

        'save commands file
        Dim oFile As New IO.StreamWriter(sFileName)
        oFile.WriteLine("rescan")
        oFile.WriteLine("select disk " + DiskNum.ToString)
        oFile.WriteLine("attributes disk clear readonly")
        oFile.WriteLine("create partition primary")
        oFile.Flush()
        oFile.Close()

        'produce the command line string
        Dim sCmd As String = "diskpart /S """ + sFileName + """" ' > """ + sLog + """"
        'Dim sCmd As String = "diskpart /S c:\1.tmp > c:\2.tmp"

        'execute
        Dim iRes As Integer = Shell(sCmd, AppWinStyle.Hide, True, 60000)

        'return results
        If iRes <> 0 Then
            EnviroBuilder.WriteToLog("Diskpart CREATE PAETITION failed", EventLogEntryType.Warning)
            Return False
        End If

        EnviroBuilder.WriteToLog("Created partition for disk " + DiskNum.ToString, EventLogEntryType.Information)
        Return True

    End Function


    ''' <summary>
    ''' Executes diskpart to a letter to the specified disk. Returns True on success.
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function DiskPartLetter(ByVal DiskNum As Integer, ByVal Letter As Char) As Boolean

        'generate file names
        Dim sFileName As String = IO.Path.GetTempFileName
        Dim sLog As String = IO.Path.GetTempFileName

        'save commands file
        Dim oFile As New IO.StreamWriter(sFileName)
        oFile.WriteLine("rescan")
        'oFile.WriteLine("select disk " + DiskNum.ToString)
        oFile.WriteLine("select volume " + DiskNum.ToString)
        oFile.WriteLine("assign letter=" + Letter)
        oFile.Flush()
        oFile.Close()

        'produce the command line string
        Dim sCmd As String = "diskpart /S """ + sFileName + """" ' > """ + sLog + """"

        'execute
        Dim iRes As Integer = Shell(sCmd, AppWinStyle.Hide, True, 60000)

        'return results
        If iRes <> 0 Then
            EnviroBuilder.WriteToLog("Diskpart ASSIGN failed: " + Letter + " to disk " + DiskNum.ToString, EventLogEntryType.Warning)
            Return False
        End If

        EnviroBuilder.WriteToLog("Assigned letter " + Letter + " to disk " + DiskNum.ToString, EventLogEntryType.Information)

        Return True

    End Function





#End Region


#Region " STATUS CHECKS "

    ''' <summary>
    ''' Returns attachment status for the specified volume
    ''' </summary>
    ''' <param name="VolumeID"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function UpdateAttachmentStatus(ByVal VolumeID As String) As String

        'prepare a request
        Dim oClient As New Amazon.EC2.AmazonEC2Client(EnviroBuilder.sAwsKey, EnviroBuilder.sAwsKeySecret)
        Dim oVolReq As New Amazon.EC2.Model.DescribeVolumesRequest
        oVolReq.VolumeId.Add(VolumeID)

        'get volume description
        Dim oVolResp As Amazon.EC2.Model.DescribeVolumesResponse = Nothing
        oVolResp = oClient.DescribeVolumes(oVolReq)

        'check status and instance ID
        Dim oVolume As Amazon.EC2.Model.Volume = oVolResp.DescribeVolumesResult.Volume(0)

        For Each oAttachment As Amazon.EC2.Model.Attachment In oVolume.Attachment
            Return oAttachment.Status
        Next

        'not attached
        Return Nothing

    End Function

    ''' <summary>
    ''' Checks volume status
    ''' </summary>
    ''' <param name="VolumeID"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function UpdateVolumeStatus(ByVal VolumeID As String) As String

        'prepare a request
        Dim oClient As New Amazon.EC2.AmazonEC2Client(EnviroBuilder.sAwsKey, EnviroBuilder.sAwsKeySecret)
        Dim oVolReq As New Amazon.EC2.Model.DescribeVolumesRequest
        oVolReq.VolumeId.Add(VolumeID)

        'get volume description
        Dim oVolResp As Amazon.EC2.Model.DescribeVolumesResponse = Nothing
        oVolResp = oClient.DescribeVolumes(oVolReq)

        If oVolResp.DescribeVolumesResult.Volume.Count = 0 Then Return Nothing

        'check status and instance ID
        Dim oVolume As Amazon.EC2.Model.Volume = oVolResp.DescribeVolumesResult.Volume(0)

        Return oVolume.Status


    End Function



    ''' <summary>
    ''' Volume status in relation to a particular instance
    ''' </summary>
    ''' <remarks></remarks>
    Private Enum enAttachmentStatus

        Attached = 0
        CanAttach = 1
        CannotAttach = 2

    End Enum


    ''' <summary>
    ''' Returns volume status in relation to the current instance
    ''' </summary>
    ''' <param name="VolumeID"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function VolumeStatus(ByVal VolumeID As String) As enAttachmentStatus

        'prepare a request
        Dim oClient As New Amazon.EC2.AmazonEC2Client(EnviroBuilder.sAwsKey, EnviroBuilder.sAwsKeySecret)
        Dim oVolReq As New Amazon.EC2.Model.DescribeVolumesRequest
        oVolReq.VolumeId.Add(VolumeID)

        'get volume description
        Dim oVolResp As Amazon.EC2.Model.DescribeVolumesResponse = Nothing
        oVolResp = oClient.DescribeVolumes(oVolReq)

        'do we have a volume at all?
        If oVolResp.DescribeVolumesResult.Volume.Count <> 1 Then
            EnviroBuilder.WriteToLog("EBS not found for attachment: " + VolumeID, EventLogEntryType.Error)
            Return enAttachmentStatus.CannotAttach
        End If

        'check status and instance ID
        Dim oVolume As Amazon.EC2.Model.Volume = oVolResp.DescribeVolumesResult.Volume(0)

        Select Case oVolume.Status
            Case "deleting" : Return enAttachmentStatus.CannotAttach

            Case "creating"
                'got to wait for this volume
                Dim dWaitFrom As DateTime = DateTime.Now
                For i As Integer = 0 To 20
                    Dim sStatUpd As String = UpdateVolumeStatus(VolumeID)
                    Select Case sStatUpd
                        Case "creating"
                            'still creating
                            Threading.Thread.CurrentThread.Sleep(10000)
                            Continue For
                        Case "available"
                            Return enAttachmentStatus.CanAttach
                        Case Else
                            EnviroBuilder.WriteToLog("Status for volume " + VolumeID + " is set to " + sStatUpd + " after creation", EventLogEntryType.Warning)
                            Return enAttachmentStatus.CannotAttach
                    End Select
                Next

            Case "in-use"
                'is it attached to this instance?
                For Each oAttachment As Amazon.EC2.Model.Attachment In oVolume.Attachment
                    Dim sAttInst As String = oAttachment.InstanceId
                    If sAttInst = EnviroBuilder.sInstanceID Then Return enAttachmentStatus.Attached
                    'attached somewhere else
                    Return enAttachmentStatus.CannotAttach
                Next

            Case "available"
                'no attachement found
                Return enAttachmentStatus.CanAttach
        End Select

        'something new 
        EnviroBuilder.WriteToLog("Unknown status for volume " + VolumeID + ": " + oVolume.Status, EventLogEntryType.Warning)
        Return enAttachmentStatus.CannotAttach

    End Function

#End Region



#Region " WMI "


    ''' <summary>
    ''' Creates a new partition for the entire disk where no partitions found
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Shared Function CreatePartitions() As Boolean

        For Each oD As WmiDisk In EnviroBuilder.arDisksPhys
            Try
                If oD.Partitions = 0 AndAlso oD.IsHDD Then DiskPartPartition(oD.Index)
            Catch ex As Exception
                EnviroBuilder.WriteToLog("Cannot create a partition." + vbCrLf + ex.Message, EventLogEntryType.Error)
                Return False
            End Try
        Next

        Return True

    End Function


    '''' <summary>
    '''' Returns true if a letter was assigned or already existed
    '''' </summary>
    '''' <param name="Letter"></param>
    '''' <returns></returns>
    '''' <remarks></remarks>
    'Private Shared Function AssignLetter(ByVal Letter As Char) As Boolean

    '    'do we already have a drive like that?
    '    Dim arDriveNames() As String = Environment.GetLogicalDrives
    '    For Each sDriveName As String In arDriveNames
    '        If sDriveName.ToLower.StartsWith(Letter) Then Return True
    '    Next

    '    'find the first volume with no drive name




    'End Function





#End Region



End Class
